import socket
import random
import struct
import time
import threading

# 统一的包格式 (共10字节头部 + 数据)
# 'B' - 类型 (1字节)
# 'H' - 序列号 (2字节)
# 'H' - 确认号 (2字节)
# 'f' - 时间戳 (4字节)
# 'B' - 数据长度 (1字节)

# 1:SYN
# 2:SYN-ACK报文
# 3：数据报文。
# 4:ACK
# 5：报文类型为 FIN，用于标识关闭连接请求
HEADER_FORMAT = 'BHHfB'
HEADER_SIZE = struct.calcsize(HEADER_FORMAT)  # 10字节

# 定义类型名称映射
TYPE_NAMES = {
    1: "SYN",
    2: "SYN-ACK",
    3: "DATA",
    4: "ACK",
    5: "FIN"
}


class ClientHandler(threading.Thread):
    """客户端处理线程：每个客户端由独立线程处理"""
    def __init__(self, sock, addr, drop_rate):
        super().__init__()
        self.sock = sock
        self.addr = addr
        self.drop_rate = drop_rate
        self.expected_seq = 0
        self.running = True

    def run(self):
        """处理客户端的所有通信（连接、数据传输、关闭）"""
        print(f"新线程处理客户端 {self.addr}")
        while self.running:
            try:
                data, addr = self.sock.recvfrom(1024)
                if addr != self.addr:
                    continue  # 只处理当前客户端的数据包
                if len(data) < HEADER_SIZE:
                    continue  # 忽略不完整包

                ptype = struct.unpack('B', data[:1])[0]

                if ptype == 1:  # SYN
                    print(f"客户端{self.addr}：收到SYN报文")
                    syn_ack = self._make_packet(2, 0, 1, b'')
                    self.sock.sendto(syn_ack, self.addr)
                    self.expected_seq = 0

                elif ptype == 3:  # 数据报文
                    header = data[:HEADER_SIZE]
                    ptype, seq, ack, timestamp, length = struct.unpack(HEADER_FORMAT, header)
                    if len(data) < HEADER_SIZE + length:
                        print(f"客户端{self.addr}：收到不完整数据包{seq}，丢弃")
                        continue

                    # 解析数据内容
                    #errors='replace' 表示当遇到无法解码的字符时，会用 ? 替换这些字符，避免因解码错误而抛出异常。
                    packet_info = data[HEADER_SIZE:HEADER_SIZE+length].decode('utf-8', errors='replace')
                    print(f"客户端{self.addr} 内容: {packet_info}")

                    # 模拟丢包
                    if random.random() < self.drop_rate:
                        print(f"客户端{self.addr}：模拟丢包 序列号{seq}")
                        continue

                    # 处理按序/失序数据包
                    if seq == self.expected_seq: #检查当前接收到的数据包的序列号 seq 是否等于服务器期望接收的序列号
                        print(f"客户端{self.addr}：正确收到数据包{seq}")
                        self.expected_seq += 1
                        ack_packet = self._make_packet(4, 0, self.expected_seq - 1, b'')
                        self.sock.sendto(ack_packet, self.addr)
                        print(f"客户端{self.addr}：发送ACK{seq}")
                    else:
                        print(f"客户端{self.addr}：失序数据包{seq}（期望{self.expected_seq}），发送ACK{self.expected_seq-1}")
                        ack_packet = self._make_packet(4, 0, self.expected_seq-1, b'')
                        self.sock.sendto(ack_packet, self.addr)

                elif ptype == 4:  # ACK
                    header = data[:HEADER_SIZE]
                    ptype, seq, ack, timestamp, length = struct.unpack(HEADER_FORMAT, header)
                    print(f"客户端{self.addr}：收到ACK报文，确认号为 {ack}")

                elif ptype == 5:  # FIN
                    print(f"客户端{self.addr}：收到FIN报文")
                    # 发送ACK确认
                    ack = self._make_packet(4, 0, 0, b'')
                    self.sock.sendto(ack, self.addr)
                    time.sleep(0.1)  # 确保ACK送达
                    # 发送FIN
                    fin = self._make_packet(5, 0, 0, b'')
                    self.sock.sendto(fin, self.addr)
                    self.running = False  # 结束线程

            except Exception as e:
                print(f"客户端{self.addr} 错误: {e}")
                self.running = False

        print(f"客户端{self.addr}：连接关闭，线程结束")

    def _make_packet(self, type, seq, ack, data):
        """构造协议报文"""
        time_send = time.time()
        header = struct.pack(HEADER_FORMAT, type, seq, ack, time_send, len(data))
        type_name = TYPE_NAMES.get(type, "UNKNOWN") # "UNKNOWN" 是一个可选参数，表示如果指定的键在字典中不存在时返回的默认值。
        packet_info = f"{type_name}-{seq}-{ack}-{time_send}-{data.decode('utf-8', errors='replace')}"
        return header + packet_info.encode('utf-8')


class ReliableUDPServer:
    def __init__(self, port, drop_rate=0.2):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.bind(('0.0.0.0', port))
        self.drop_rate = drop_rate
        self.clients = {}  # 客户端地址 -> 处理线程

    def start(self):
        print(f"服务器启动，监听端口{self.sock.getsockname()[1]}，丢包率{self.drop_rate*100}%")
        try:
            while True:
                data, addr = self.sock.recvfrom(1024)
                if addr not in self.clients or not self.clients[addr].is_alive():
                    # 新客户端，创建处理线程
                    handler = ClientHandler(self.sock, addr, self.drop_rate)
                    self.clients[addr] = handler #将新创建的处理线程添加到 self.clients 字典中
                    handler.start()
        except KeyboardInterrupt: #用户手动终止服务器的异常
            print("服务器关闭")
        finally:
            self.sock.close()


if __name__ == "__main__":
    server_port = int(input("请输入服务器端口号: "))
    drop_rate = float(input("设置丢包率（0-1，默认0.2）: ") or "0.2")
    drop_rate = max(0.0, min(1.0, drop_rate))  # 限制范围
    server = ReliableUDPServer(server_port, drop_rate)
    server.start()